#!/usr/bin/env python3
import argparse
import os
import sys
from time import sleep
from time import strftime
import grpc
from scapy.all import Ether, Packet, BitField, raw
import ipaddress

# Import P4Runtime lib from parent utils dir
# Probably there's a better way of doing this.
sys.path.append(
    os.path.join(os.path.dirname(os.path.abspath(__file__)),
                 '../../utils/'))
import p4runtime_lib.bmv2
import p4runtime_lib.helper
from p4runtime_lib.switch import ShutdownAllSwitchConnections

class CpuHeader(Packet):
    name = 'CpuPacket'
    fields_desc = [
        BitField('src_ip', 0, 32),
        BitField('dst_ip', 0, 32),
        BitField('egress_port', 0, 9),
        BitField('ecn_threshold', 0, 19),
        BitField('pad', 0, 4)
    ]

def writeCloneSessionEntry(p4info_helper, sw, clone_session_id, cpu_port = 10):
    """
    Installs a clone session entry to a switch

    :param p4info_helper: the P4Info helper
    :param sw: the switch connection
    :param clone_session_id: the session identifier
    :param cpu_port: the CPU port of the switch, default to 10
    """
    clone_entry = p4info_helper.buildCloneSessionEntry(clone_session_id,
                                                       [{"egress_port": cpu_port, "instance": 0}], 0)
    sw.WritePREEntry(clone_entry)

def writeTableRule(p4info_helper, sw, table_name, match_fields, action_name, action_params = None):
    """
    Installs a table rule to a switch

    :param p4info_helper: the P4Info helper
    :param sw: the switch connection
    :param table_name: the match table name
    :param match_fields: match fields dictionary
    :param action_name: action name
    :param action_params: the parameters of the action, default to None
    """
    table_entry = p4info_helper.buildTableEntry(
        table_name=table_name,
        match_fields=match_fields,
        action_name=action_name,
        action_params=action_params)
    sw.WriteTableEntry(table_entry)

def writeIpv4ForwardRule(p4info_helper, sw, dst_ip_addr, forward_mac_addr, forward_port, dst_ip_length=32):
    writeTableRule(p4info_helper, sw, 
        table_name="MyIngress.ipv4_lpm", 
        match_fields={
            "hdr.ipv4.dstAddr": (dst_ip_addr, dst_ip_length)
        },
        action_name="MyIngress.ipv4_forward",
        action_params={
            "dstAddr": forward_mac_addr,
            "port": forward_port
        }
    )


def readTableRules(p4info_helper, sw):
    """
    Reads the table entries from all tables on the switch.

    :param p4info_helper: the P4Info helper
    :param sw: the switch connection
    """
    print('\n----- Reading tables rules for %s -----' % sw.name)
    for response in sw.ReadTableEntries():
        for entity in response.entities:
            entry = entity.table_entry
            table_name = p4info_helper.get_tables_name(entry.table_id)
            print('%s: ' % table_name, end=' ')
            for m in entry.match:
                print(p4info_helper.get_match_field_name(table_name, m.field_id), end=' ')
                print('%r' % (p4info_helper.get_match_field_value(m),), end=' ')
            action = entry.action.action
            action_name = p4info_helper.get_actions_name(action.action_id)
            print('->', action_name, end=' ')
            for p in action.params:
                print(p4info_helper.get_action_param_name(action_name, p.param_id), end=' ')
                print('%r' % p.value, end=' ')
            print()

def printGrpcError(e):
    print("gRPC Error:", e.details(), end=' ')
    status_code = e.code()
    print("(%s)" % status_code.name, end=' ')
    traceback = sys.exc_info()[2]
    print("[%s:%d]" % (traceback.tb_frame.f_code.co_filename, traceback.tb_lineno))

def main(p4info_file_path, bmv2_file_path):
    # Instantiate a P4Runtime helper from the p4info file
    p4info_helper = p4runtime_lib.helper.P4InfoHelper(p4info_file_path)

    try:
        print("please input the length of ecn: ", end="")
        ecn_length = int(input())

        # Create a switch connection object for s1, s2 and s3;
        # this is backed by a P4Runtime gRPC connection.
        # Also, dump all P4Runtime messages sent to switch to given txt files.
        s1 = p4runtime_lib.bmv2.Bmv2SwitchConnection(
            name='s1',
            address='127.0.0.1:50051',
            device_id=0,
            proto_dump_file='logs/s1-p4runtime-requests.txt')
        s2 = p4runtime_lib.bmv2.Bmv2SwitchConnection(
            name='s2',
            address='127.0.0.1:50052',
            device_id=1,
            proto_dump_file='logs/s2-p4runtime-requests.txt')
        s3 = p4runtime_lib.bmv2.Bmv2SwitchConnection(
            name='s3',
            address='127.0.0.1:50053',
            device_id=2,
            proto_dump_file='logs/s3-p4runtime-requests.txt')

        # Send master arbitration update message to establish this controller as
        # master (required by P4Runtime before performing any other write operation)
        s1.MasterArbitrationUpdate()
        s2.MasterArbitrationUpdate()
        s3.MasterArbitrationUpdate()

        # Install the P4 program on the switches
        s1.SetForwardingPipelineConfig(p4info=p4info_helper.p4info,
                                       bmv2_json_file_path=bmv2_file_path)
        print("Installed P4 Program using SetForwardingPipelineConfig on s1")
        s2.SetForwardingPipelineConfig(p4info=p4info_helper.p4info,
                                       bmv2_json_file_path=bmv2_file_path)
        print("Installed P4 Program using SetForwardingPipelineConfig on s2")
        s3.SetForwardingPipelineConfig(p4info=p4info_helper.p4info,
                                       bmv2_json_file_path=bmv2_file_path)
        print("Installed P4 Program using SetForwardingPipelineConfig on s3")

        # Install necessary ipv4 forward rules
        # s1 rules
        writeIpv4ForwardRule(p4info_helper, s1, 
            dst_ip_addr="10.0.1.1",
            forward_mac_addr="08:00:00:00:01:01",
            forward_port=2)
        writeIpv4ForwardRule(p4info_helper, s1, 
            dst_ip_addr="10.0.1.11",
            forward_mac_addr="08:00:00:00:01:11",
            forward_port=1)
        writeIpv4ForwardRule(p4info_helper, s1, 
            dst_ip_addr="10.0.2.0",
            dst_ip_length=24,
            forward_mac_addr="08:00:00:00:02:00",
            forward_port=3)
        writeIpv4ForwardRule(p4info_helper, s1, 
            dst_ip_addr="10.0.3.0",
            dst_ip_length=24,
            forward_mac_addr="08:00:00:00:03:00",
            forward_port=4)
        # s2 rules
        writeIpv4ForwardRule(p4info_helper, s2, 
            dst_ip_addr="10.0.2.2",
            forward_mac_addr="08:00:00:00:02:02",
            forward_port=2)
        writeIpv4ForwardRule(p4info_helper, s2, 
            dst_ip_addr="10.0.2.22",
            forward_mac_addr="08:00:00:00:02:22",
            forward_port=1)
        writeIpv4ForwardRule(p4info_helper, s2, 
            dst_ip_addr="10.0.1.0",
            dst_ip_length=24,
            forward_mac_addr="08:00:00:00:01:00",
            forward_port=3)
        writeIpv4ForwardRule(p4info_helper, s2, 
            dst_ip_addr="10.0.3.0",
            dst_ip_length=24,
            forward_mac_addr="08:00:00:00:03:00",
            forward_port=4)
        # s3 rules
        writeIpv4ForwardRule(p4info_helper, s3, 
            dst_ip_addr="10.0.3.3",
            forward_mac_addr="08:00:00:00:03:03",
            forward_port=1)
        writeIpv4ForwardRule(p4info_helper, s3, 
            dst_ip_addr="10.0.1.0",
            dst_ip_length=24,
            forward_mac_addr="08:00:00:00:01:00",
            forward_port=2)
        writeIpv4ForwardRule(p4info_helper, s3, 
            dst_ip_addr="10.0.2.0",
            dst_ip_length=24,
            forward_mac_addr="08:00:00:00:02:00",
            forward_port=3)
        print("Installed IPv4 forward rules")
        
        # Install ecn rules
        writeTableRule(p4info_helper, s1, 
            table_name="MyEgress.ecn_queue_threshold", 
            match_fields={
                "standard_metadata.egress_port": 3
            },
            action_name="MyEgress.set_ecn_threshold",
            action_params={
                "threshold": ecn_length
            }
        )
        writeTableRule(p4info_helper, s2, 
            table_name="MyEgress.ecn_queue_threshold", 
            match_fields={
                "standard_metadata.egress_port": 3
            },
            action_name="MyEgress.set_ecn_threshold",
            action_params={
                "threshold": ecn_length
            }
        )
        print("Installed ecn threshold rules")

        # Install clone session entry
        writeCloneSessionEntry(p4info_helper, s1, 100)
        writeCloneSessionEntry(p4info_helper, s2, 100)
        print("Installed clone session entries")

        print("Waiting for reports...")
        while True:
            for msg in s1.stream_msg_resp:
                # Only process PacketIn
                if (msg.WhichOneof('update') == 'packet'):
                    packet = Ether(raw(msg.packet.payload))
                    # Parse CPU header, notify user about the congestion
                    if packet.type == 0x2333:
                        cpu_header = CpuHeader(bytes(packet.load))
                        print("S1报告发生了链路拥塞，下面是发生拥塞时S1转发的数据包的相关信息")
                        print(f"源IP地址：{str(ipaddress.IPv4Address(cpu_header.src_ip))}")
                        print(f"目的IP地址：{str(ipaddress.IPv4Address(cpu_header.dst_ip))}")
                        print(f"出端口：{cpu_header.egress_port}")
                        print(f"当前队列长度阈值：{cpu_header.ecn_threshold}")
            for msg in s2.stream_msg_resp:
                # Only process PacketIn
                if (msg.WhichOneof('update') == 'packet'):
                    packet = Ether(raw(msg.packet.payload))
                    # Parse CPU header, notify user about the congestion
                    if packet.type == 0x2333:
                        cpu_header = CpuHeader(bytes(packet.load))
                        print("S2报告发生了链路拥塞，下面是发生拥塞时S2转发的数据包的相关信息")
                        print(f"源IP地址：{str(ipaddress.IPv4Address(cpu_header.src_ip))}")
                        print(f"目的IP地址：{str(ipaddress.IPv4Address(cpu_header.dst_ip))}")
                        print(f"出端口：{cpu_header.egress_port}")
                        print(f"当前队列长度阈值：{cpu_header.ecn_threshold}")
            sleep(1)

        # Read table entries from s1, s2 and s3
        # readTableRules(p4info_helper, s1)
        # readTableRules(p4info_helper, s2)
        # readTableRules(p4info_helper, s3)
    except KeyboardInterrupt:
        print(" Shutting down.")
    except grpc.RpcError as e:
        printGrpcError(e)

    ShutdownAllSwitchConnections()

if __name__ == '__main__':
    parser = argparse.ArgumentParser(description='P4Runtime Controller')
    parser.add_argument('--p4info', help='p4info proto in text format from p4c',
                        type=str, action="store", required=False,
                        default='./build/ecn.p4.p4info.txt')
    parser.add_argument('--bmv2-json', help='BMv2 JSON file from p4c',
                        type=str, action="store", required=False,
                        default='./build/ecn.json')
    args = parser.parse_args()

    if not os.path.exists(args.p4info):
        parser.print_help()
        print("\np4info file not found: %s\nHave you run 'make'?" % args.p4info)
        parser.exit(1)
    if not os.path.exists(args.bmv2_json):
        parser.print_help()
        print("\nBMv2 JSON file not found: %s\nHave you run 'make'?" % args.bmv2_json)
        parser.exit(1)
    main(args.p4info, args.bmv2_json)
